#pragma once

#include <cmath>

#include "base.h"
#include "../impl/filter.h"
#include "../impl/ma.h"
#include "../updater/count.h"
#include "../updater/ohlc.h"

namespace mdtx::online::sampler
{
    /**
     * @brief Imbalance bar
     *
     * @cite De Prado, M. L. (2018). Advances in financial machine learning. John Wiley & Sons.
     *
     * @tparam Time
     * @tparam T
     * @tparam Idx index of input tuple to calculate imbalance, 0 for tick
     */
    template <typename Time, typename T, size_t Idx>
    class imbalance : public sampler_base<Time, T, 3, 7>
    {
        using base = sampler_base<Time, T, 3, 7>;

    public:
        using typename base::tuple_in;
        using typename base::tuple_out;

    protected:
        updater::ohlc_tumbling<T> ohlc{};
        updater::vwap<T> vwap{};
        ptrdiff_t ntick{}, imba_dir{}, warmup;

        Time now{};
        T last{}, imba{}, thres{};

        impl::sum<T> imba_cusum{};
        impl::ema<T> ntick_exp, imba_exp;
        bool init{};

    public:
        /**
         * @brief Construct a new imbalance object
         *
         * @param alpha smoothing factor for EMA
         * @param warmup number of ticks to warmup initial sampling threshold
         */
        imbalance(T alpha, ptrdiff_t warmup)
            : warmup(warmup),
              ntick_exp(alpha),
              imba_exp(alpha)
        {
        }

        tuple_out update(tuple_in const &in) override
        {
            auto [time, price, vol, tnvr] = in;
            now = time;

            if (!init)
            {
                last = price;
                init = true;
            }

            tuple_out sample{};
            ohlc.insert_impl(price);
            vwap.insert_impl(vol, tnvr);
            ntick += 1;

            auto dir = sign(price - last);
            if (dir)
            {
                imba_dir = dir;
            }
            if constexpr (Idx == 0)
            {
                // -> tick
                imba = imba_dir;
            }
            else
            {
                imba = imba_dir * std::get<Idx>(in);
            }
            imba_cusum.insert_impl(imba);
            imba_exp.insert_impl(imba);

            if (warmup)
            {
                warmup -= 1;
                if (!warmup)
                {
                    // run once, initialise threshold after warmup
                    thres = std::abs(imba_exp.m) * ntick;
                }
            }
            // TODO: we should extensively test the behaviour of sliding threshold and tumbling threshold
            else if (std::abs(imba_cusum.s) > thres)
            {
                sample = detail::cat_time_tups(time, ohlc.pop(), vwap.pop());
                ntick_exp.insert_impl(std::exchange(ntick, 0));
                imba_cusum.reset_impl();
                thres = std::abs(imba_exp.m) * ntick_exp.m;
            }

            return sample;
        }

        tuple_out flush() override
        {
            tuple_out sample{};
            if (ohlc.init)
            {
                sample = detail::cat_time_tups(now, ohlc.pop(), vwap.pop());
            }
            return sample;
        }
    };

    /**
     * @brief Sample on cumsum of log return
     *
     * @tparam Time
     * @tparam T
     * @tparam Idx index of input tuple to calculate imbalance, 0 equivalent to tumbling time window
     */
    template <typename Time, typename T, size_t Idx>
    class cusum_logret : public sampler_base<Time, T, 3, 7>
    {
        using base = sampler_base<Time, T, 3, 7>;

    public:
        using typename base::tuple_in;
        using typename base::tuple_out;

    protected:
        updater::ohlc_tumbling<T> ohlc{};
        updater::vwap<T> vwap{};

        Time now{};
        T last{};

        impl::cusum_filter<T> cusum;
        bool init{};

    public:
        explicit cusum_logret(T ret_threshold)
            : cusum(ret_threshold)
        {
        }

        tuple_out update(tuple_in const &in) override
        {
            auto [time, price, vol, tnvr] = in;
            auto fval = std::log(std::get<Idx>(in)); // filtered value
            now = time;

            tuple_out sample{};
            ohlc.insert_impl(price);
            vwap.insert_impl(vol, tnvr);

            if (!init)
            {
                last = fval;
                init = true;
            }

            cusum.insert_impl(fval - last); // no drift
            last = fval;
            if (cusum.state)
            {
                sample = detail::cat_time_tups(now, ohlc.pop(), vwap.pop());
            }

            return sample;
        }

        tuple_out flush() override
        {
            tuple_out sample{};
            if (ohlc.init)
            {
                sample = detail::cat_time_tups(now, ohlc.pop(), vwap.pop());
            }
            return sample;
        }
    };
}
